#!/usr/bin/env python
#
# Copyright 2013 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Module for handling interactions between the testing device.

Functions in this script are used to communicate with a device running OpenWRT
and are used to collect and parse information from the system.
"""
__author__ = "Joey Freeland (joeyfreeland@gmail.com)"

import json
import logging
import os
import random
import re
import shlex
import signal
import subprocess
import time

import ref


def Associate(ssid):
  """Attempt to associate the device with a Wi-Fi network.

  Args:
    ssid: The SSID that the device should be associated with.

  Returns:
    None: Should probably look at the detailed output from this command again.
  """
  logging.info("Using iwconfig to associate with %s.", ssid)
  # TODO(joeyfreeland): Move commands elsewhere and make them platform
  # independent.
  connect_cmd = "/usr/bin/sudo /sbin/iwconfig wlan0 essid \"%s\"" % ssid
  errors = RunCommand(connect_cmd, results=False, errors=True)
  if errors:
    logging.debug("Problems with iwconfig:\n\n%s\n\n", errors)
  time.sleep(5)


def CheckAssociation(ssid):
  """Verify the device did associate with the intended SSID.

  Args:
    ssid: The SSID that the device should be associated with.

  Returns:
    Pass (1 = associated) or fail (0 = not associated).
  """
  logging.info("Confirming that we have associated with %s.", ssid)
  results, errors = RunCommand("/usr/bin/sudo /sbin/iwconfig wlan0",
                               errors=True)
  if errors:
    logging.debug("Problems with association:\n%s", errors)
  wlan_search = re.search(r"wlan0.*ESSID\:\"(.*)\"", results)
  if wlan_search:
    logging.info("Successfully associated with %s.", ssid)
    return 1
  else:
    logging.debug("Did not associate with %s.", ssid)
    return 0


def CheckCurrentWiFiSpecs():
  """Have the device check the WiFi signal and speed using 'iw'.

  Args:
    None

  Returns:
    my_signal: The signal strength in dBm.  Should always be a negative number.
    my_speed: The data rate which was established with the access point.
  """
  results = RunCommand("/usr/bin/sudo /sbin/iwconfig wlan0")
  # TODO(joeyfreeland): This regex will need to be extended for dual-band
  # support.
  frequency = re.search(r"Frequency.*(2.4\d{2})", results).group(1)
  my_channel = ref.WIFI_CHANNEL_MAP[frequency]
  my_speed = re.search(r"Bit Rate(\:|=)(.*) Mb\/s", results).group(2)
  signal = re.search(r"Signal level=(-\d+) dBm|Signal level=(\d+)\/100",
                     results)
  if signal.group(1) is None:
    my_signal = int(signal.group(2))
  else:
    my_signal = int(signal.group(1))
  return my_channel, my_speed, my_signal


def CheckLock():
  """Check to see if there is another process already running."""
  try:
    lock = open("lock", "r")
  except IOError:
    logging.info("Lock file does not exist.  Creating one and recording PID.")
    lock = open("lock", "w+")
    lock.write(str(os.getpid()))
    lock.close()
    return 0

  pid = lock.read()
  lock.close()
  if pid and os.path.exists("/proc/" + pid):
    logging.debug("Another process is already running.")
    return 1
  else:
    logging.debug("No other process is running.  Continuing.")
    lock = open("lock", "w+")
    lock.write(str(os.getpid()))
    lock.close()
    return 0

def ConfigureWPA(networks="", ssid=""):
  """Generate a configuration file for wpa_supplicant.

  Check to see if the network requires authentication and if so, assign the
  correct security settings.

  Args:
    networks: A dictionary representing all recognized networks.
    ssid: SSID of the network to configure.

  Returns:
    None
  """
  logging.info("Generating wpa_supplicant config for %s.", ssid)
  wpa_config_file = open("/home/pi/wpa_supplicant-wlan0.conf", "w+")
  wpa_config = ("ctrl_interface=/home/pi/wpa_supplicant-wlan0\n"
                "eapol_version=1\n"
                "fast_reauth=1\n"
                "network={\n"
                "ssid=\"%s\"\n" % ssid)
  if "8021x" in networks[ssid].keys():
    wpa_config += ("scan_ssid=0\n"
                   "key_mgmt=WPA-EAP\n"
                   "auth_alg=OPEN\n"
                   "eap=PEAP\n"
                   "phase1=\"peaplabel=0\"\n"
                   "phase2=\"auth=MSCHAPV2\"\n"
                   "ca_cert=\"%s\"\n"
                   "ca_path=\"/etc/ssl/certs/\"\n"
                   "identity=\"%s\"\n"
                   "password=\"%s\"\n"
                   "}\n" % (networks[ssid]["ca_cert"],
                            networks[ssid]["username"],
                            networks[ssid]["password"]))
    logging.info("Configured %s using PEAP/MSCHAPV2.", ssid)
  elif "auth" in networks[ssid].keys():
    wpa_config += "\tpsk=%s\n}\n" % networks[ssid]["wpa"]
    logging.info("Configured %s using WPA-PSK.", ssid)
  else:
    wpa_config += "\tkey_mgmt=NONE\n}"
    logging.info("Configured %s to use no authentication.", ssid)
  wpa_config_file.write(wpa_config)
  wpa_config_file.close()

"""
ctrl_interface=/home/pi/wpa_supplicant-wlan0
eapol_version=1
fast_reauth=1
network={
ssid="Wireless@SGx"
scan_ssid=0
key_mgmt=WPA-EAP
auth_alg=OPEN
eap=PEAP
phase1="peaplabel=0"
phase2="auth=MSCHAPV2"
ca_cert="/etc/ssl/certs/Go_Daddy_Class_2_CA.pem"
ca_path="/etc/ssl/certs/"
identity="97772198@stm"
password="PP7MkH4s"
"""

def CountNeighbors():
  """Count the number of nodes in the same broadcast domain.

  Args:
    None

  Returns:
    An integer representing the number of nodes attached to this AP.
  """
  invalid_ips = (1, 252, 253, 254, 255)
  count = 0
  ip, netmask = GetIPConfig(netmask=True)
  try:
    net_size = int(ref.SUBNET[netmask])
  except:
    net_size = 0
  if net_size > 19:
    cmd = ("/usr/bin/sudo /usr/bin/arp-scan -q --interface=wlan0 %s/%d" %
           (ip, net_size))
    results = RunCommand(cmd, timeout=30)
    for line in results.split("\n"):
      try:
        last_octet = re.search(r"([0-9]+(?:\.[0-9]+){2}\.([0-9]+))",
                               line).group(2)
        if int(last_octet) not in invalid_ips: count += 1
      except:
        pass
    logging.info("Found %s neighbors on this AP.", count)
    return count
  else:
    logging.debug("Network too large. Could not check for neighbors.")
    return 0


def GetIPConfig(netmask=False):
  results = RunCommand("/usr/bin/sudo /sbin/ifconfig wlan0")
  ip = re.search(r"addr:([0-9]+(?:\.[0-9]+){3})", results).group(1)
  if netmask:
    netmask = re.search(r"Mask:([0-9]+(?:\.[0-9]+){3})", results).group(1)
    return ip, netmask
  return ip


def GetMACAddress():
  results = RunCommand("/usr/bin/sudo /sbin/ifconfig wlan0")
  mac = re.search(r"(([0-9A-Fa-f]{2}[:]){5}([0-9A-Fa-f]{2}))", results).group(0)
  return mac


def GetTemperature():
  results = RunCommand("/opt/vc/bin/vcgencmd measure_temp")
  temp = re.search(r"temp=(.*)'C", results).group(1)
  return temp


def KillWPA():
  """Kill the wpa_supplicant process.

  Args:
    None

  Returns:
    None
  """
  logging.info("Killing existing wpa_supplicant process.")
  RunCommand("/usr/bin/sudo /usr/bin/killall -9 wpa_supplicant", results=False,
             errors=True)


def LoadConfig(config_file, software_version):
  """Process or generate a configuration file.

  Args:
    config_file: The local filename of the configuration.
    software_version: The software version of the current script set.

  Returns:
    A dictionary that holds the configuration parameters of this device.
  """
  config = {}
  try:
    cfg_file = open(config_file)
    config = json.loads(cfg_file.read())
    cfg_file.close()
    config["software_version"] = software_version
    return config
  except:
    config["config_version"] = "0.0.0"
    config["device_name"] = "wtd-%d" % random.randrange(1, 100000)
    # TODO(joeyfreeland): Need to store default speedtest servers elsewhere.
    # We probably want to record the country code of the device by checking the
    # external IP address and geo-locating the device, then using speedtest
    # servers based on a pre-defined dictionary of servers.
    config["local_speedtest_name"] = "Biznet"
    config["local_speedtest_url"] = (
        "http://speedtest.biznetnetworks.com/speedtest/")
    config["mac_address"] = GetMACAddress()
    config["provider"] = "unknown"
    config["remote_speedtest_name"] = "SingTel"
    config["remote_speedtest_url"] = (
        "http://speedtest.singnet.com.sg/speedtest/")
    config["software_version"] = software_version
    cfg_file = open(config_file, "w+")
    cfg_file.write(json.dumps(config))
    cfg_file.close()
    return config


def RequestDHCP(ssid):
  """Kill existing dhclient processes and attempt to acquire new DHCP settings.

  Args:
    ssid: Simply used to log what network we're requesting DHCP from.  Should
          probably be handled in the main() function.

  Returns:
    Pass (1 = process executed) or fail (0 = process ended with system error).
  """
  # TODO(joeyfreeland): I could go nutso here and pull out the IP address,
  # netmask, gateway, and nameservers that were assigned.  For now just checking
  # that it was a success.
  logging.info("Killing any stagnant DHCP processes.")
  RunCommand("/usr/bin/sudo /usr/bin/killall dhclient", results=False,
             errors=False)
  logging.info("Requesting DHCP configuration from %s.", ssid)
  RunCommand("/usr/bin/sudo /sbin/dhclient wlan0", results=False, errors=False)
  time.sleep(3)


def RunCommand(cmd, kill_tree=False, timeout=10, results=True, errors=False):
  """Run a command with a timeout after which it will be forcibly killed."""
  class Alarm(Exception):
    pass
  def alarm_handler(signum, frame):
    raise Alarm
  cmd = shlex.split(cmd)
  p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                       shell=False)
  if timeout > 1:
    signal.signal(signal.SIGALRM, alarm_handler)
    signal.alarm(timeout)
  try:
    stdout, stderr = p.communicate()
    if timeout > 1:
      signal.alarm(0)
  except Alarm:
    pids = [p.pid]
    if kill_tree:
      pids.extend(GetProcessChildren(p.pid))
    for pid in pids:
      try:
        os.kill(pid, signal.SIGKILL)
      except:
        pass
    if results and not errors:
      return ""
    elif errors and not results:
      return ""
    else:
      return "", ""
  if results and not errors:
    return stdout
  elif errors and not results:
    return stderr
  else:
    return stdout, stderr


def GetProcessChildren(pid):
  cmd = "/bin/ps --no-headers -o pid --ppid %d" % pid
  cmd = shlex.split(cmd)
  output = subprocess.check_output(cmd, shell=False)
  if output:
    return [int(p) for p in output.split()]
  else:
    return


def Scan(networks):
  """Scan for wireless networks, parse the output, and record various data.

  For networks we recognize, add them to a list.

  Args:
    networks: A dictionary representing all recognized networks.

  Returns:
    None
  """
  # TODO(joeyfreeland): Need to come back and add 5ghz channels to support
  # dual-band cards.
  aps_per_channel = {"1": 0, "2": 0, "3": 0, "4": 0, "5": 0, "6": 0, "7": 0,
                     "8": 0, "9": 0, "10": 0, "11": 0, "12": 0, "13": 0}
  best_signal, signal_dbm, ssids = {}, [], []
  logging.info("Scanning for networks.")
  cmd = "/usr/bin/sudo /sbin/iwlist wlan0 scan"
  results, errors = RunCommand(cmd, errors=True)
  ifdown = re.search(r"Interface doesn't support scanning", errors)
  if ifdown:
    logging.info("Interface does not exist.  Restarting WPA.")
    return 0, 0, 0, 0, 0
  results = results.split("Cell")
  for cell in results:
    channel_search = re.search(r"Channel:(.*)|\(Channel (.*)\)", cell)
    if channel_search:
      if channel_search.group(1) is None:
        channel = channel_search.group(2)
      else:
        channel = channel_search.group(1)
      aps_per_channel[channel] += 1
    else:
      continue

    signal = re.search(r"Signal level=(.*) dBm|Signal level=(.*)\/100",
                       cell)
    if signal.group(1) is None:
      signal_dbm.append(int(signal.group(2)))
    else:
      signal_dbm.append(int(signal.group(1)))

    ssid = re.search(r"ESSID:\"(.*)\"", cell).group(1)

    if ssid not in ssids:
      ssids.append(ssid)
      if ssid in networks.keys():
        logging.info("Found a recognized network: %s", ssid)
      else:
        logging.info("Found an unknown network: %s", ssid)

    if ssid in best_signal.keys():
      if signal < best_signal[ssid]:
        best_signal[ssid] = signal
    else:
      best_signal[ssid] = signal

  return aps_per_channel, signal_dbm, ssids, best_signal


def StartWPA():
  """Start the wpa_supplicant process."""
  logging.info("Starting wpa_supplicant process.")
  cmd = ("/usr/bin/sudo /sbin/wpa_supplicant -B -qq -D nl80211,wext -i wlan0 -c"
         " /home/pi/wpa_supplicant-wlan0.conf")
  errors = RunCommand(cmd, results=False, errors=True)
  if errors:
    logging.debug("There were errors starting WPA: %s", errors)
